import os
import docx
import re

from preprocess.name_sep import extract_name_from_text

def extract_docx_text(docFile):
    document = docx.Document(docFile)
    docx_text = ""
    for para in document.paragraphs:
        docx_text += para.text + '\n'
    return docx_text

def extract_key_sections(docFile, output_file=None):
    text = extract_docx_text(docFile)
    lines = text.split('\n')

    # 用于存储提取的关键部分
    admission_record = []
    death_record = []
    death_case_discussion = []

    # 当前处理的部分
    current_section = None

    # 查找关键部分的模式
    admission_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}\s.*入院记录')
    death_record_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}\s.*死亡记录')
    death_discussion_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}\s.*死亡病例讨论')

    # 章节结束的模式（新章节开始）
    stop_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}\s.*')

    # 遍历每一行
    i = 0
    while i < len(lines):
        line = lines[i]

        # 检查是否是需要提取的章节开始
        if admission_pattern.match(line):
            current_section = "admission"
            admission_record.append(f"Line {i}: {line}")

            # 收集章节内容直到下一个日期时间格式行
            j = i + 1
            while j < len(lines) and not stop_pattern.match(lines[j]):
                if lines[j].strip():  # 只添加非空行
                    # admission_record.append(f"Line {j}: {lines[j]}")
                    admission_record.append(lines[j])
                j += 1
            i = j - 1  # 减1是因为循环结束后会自动 +1

        elif death_record_pattern.match(line):
            current_section = "death_record"
            death_record.append(f"Line {i}: {line}")

            # 收集章节内容直到下一个日期时间格式行
            j = i + 1
            while j < len(lines) and not stop_pattern.match(lines[j]):
                if lines[j].strip():  # 只添加非空行
                    # death_record.append(f"Line {j}: {lines[j]}")
                    death_record.append(lines[j])
                j += 1
            i = j - 1

        elif death_discussion_pattern.match(line):
            current_section = "death_discussion"
            death_case_discussion.append(f"Line {i}: {line}")

            # 收集章节内容直到下一个日期时间格式行
            j = i + 1
            while j < len(lines) and not stop_pattern.match(lines[j]):
                if lines[j].strip():  # 只添加非空行
                    # death_case_discussion.append(f"Line {j}: {lines[j]}")
                    death_case_discussion.append(lines[j])
                j += 1
            i = j - 1

        i += 1

    # 组织提取的部分
    results = {
        "入院记录": admission_record,
        "死亡记录": death_record,
        "死亡病例讨论": death_case_discussion
    }

    # 如果指定了输出文件，则写入文件
    if output_file:
        with open(output_file, 'w', encoding='utf-8') as f:
            for section_name, content in results.items():
                f.write(f"\n\n===== {section_name} =====\n\n")
                for line in content:
                    f.write(line + '\n')

    return results

# 示例使用
# dict_path = r'C:\Users\yyds\Desktop\new_demo'
# docFile = os.path.join(dict_path, "519038.docx")
#
# text = extract_docx_text(docFile)
# # 提取姓名
# name = extract_name_from_text(text)
# print("姓名：" + name)
#
# sections = extract_key_sections(docFile, f'519038{name}.txt')
#
# # 打印提取的部分
# for section_name, content in sections.items():
#     print(f"\n{section_name}部分找到 {len(content)} 行")
#     # 显示每个部分的前几行作为示例
#     for line in content[:5]:
#         print(line)